import cv2
import os
import numpy as np
import argparse
import math
import csv


parser = argparse.ArgumentParser(description='list to image')

parser.add_argument('--name', default='nature',
                    type=str, help='config file')
parser.add_argument('--path', default='/home/xulong/snowstorm/0Dataset/VOT2018/nature/',
                    type=str, help='config file')

args = parser.parse_args()

list0 = []
root_path = args.path

for files in os.walk(root_path):
    list0.append(files)
img_list = []
for ii in list0[0][2]:
    img_list.append(ii)
list_img_name = np.array(img_list)
list_img_name.sort()
# print(list_img_name)

# csv_name = args.name + '.csv'
csv_name = args.name + '.txt'
csv_path = os.path.join(root_path, csv_name)
csv_file = csv.reader(open(csv_path))

list2 = []
for i in csv_file:
    list2.append(i)
list3 = np.array(list2)
list4 = []
list_temp = []
for j in list3:
    list_temp = np.array([0.0 if y == '' else float(y) for y in j])
    list4.append(list_temp)
list5 = np.array(list4)
list_temp1 = []
list_temp2 = []
for k in list5:
    list_temp1 = np.array([int(z) for z in k])
    list_temp2.append(list_temp1)
list_rect = np.array(list_temp2)

img_num = len(list_rect) - 1
for num in range(0, img_num):
    if list_rect[num][0] > 2:
        img_file_name = root_path + str(list_img_name[num])
        src = cv2.imread(img_file_name)
        # cv2.imshow('aa',src)
        # cv2.waitKey(0)

        if len(list_rect[num]) == 4:
            x1 = list_rect[num][0]
            y1 = list_rect[num][1]
            w = list_rect[num][2]
            h = list_rect[num][3]
            x2 = x1 + w
            y2 = y1 + h
            cv2.rectangle(src, (x1, y1), (x2, y2), (0, 0, 255), 3)
            print("number:", num, "  len:", len(list_rect[num]), "  bbox:", (x1, y1), (x2, y2))

        elif len(list_rect[num]) == 8:
            x1 = min(list_rect[num][0], list_rect[num][2], list_rect[num][4], list_rect[num][6])
            x2 = max(list_rect[num][0], list_rect[num][2], list_rect[num][4], list_rect[num][6])
            y1 = min(list_rect[num][1], list_rect[num][3], list_rect[num][5], list_rect[num][7])
            y2 = max(list_rect[num][1], list_rect[num][3], list_rect[num][5], list_rect[num][7])
            w = x2 - x1
            h = y2 - y1
            cv2.rectangle(src, (x1, y1), (x2, y2), (0, 0, 255), 3)
            print("number:", num, "  len:", len(list_rect[num]), "  bbox:", (x1, y1), (x2, y2))

        else:
            print("data error!")

        # cv2.imshow('aa',src)
        # cv2.waitKey(0)

        # save_name = args.name + '_result/'
        save_name = args.name + '/'
        save_path = os.path.join(root_path, save_name)
        # save_path = os.path.join(root_path, 'result/')
        save_file_name = save_path + str(list_img_name[num])
        if not os.path.isdir(save_path):
            os.makedirs(save_path)
        cv2.imwrite(save_file_name, src)

print('Save image successfully!')

